#ifndef _CANARY2_SYS_THREAD_POOL_H_
#define _CANARY2_SYS_THREAD_POOL_H_

#include <thread>
#include <vector>
#include <functional>
#include <stop_token>
#include <coroutine>
#include <concepts>
#include <mutex>
#include "resumable.h"
#include "scheduler_concepts.h"
#include "task_manager.h"
#include "sleeper_concepts.h"
#include <iostream>

#include "../util/concurrentqueue.h"

namespace wilson::sys
{
    struct empty_invocable
    {
        constexpr void operator()() const noexcept {}
    };

    class empty_sleeper : public empty_invocable
    {
    public:
        constexpr void awake() const noexcept {}
        constexpr bool done() const noexcept { return true; }
    };

    template<task_manager       TaskM   = normal_task_mgr,
             sleeper            Sleeper = empty_sleeper, 
             ::std::invocable   Before_Exec_Task_Func = empty_invocable>
    class thread_pool
    {
    public:
        using task_manager_type = TaskM;
        using sleeper_type = Sleeper;

    public:
        thread_pool(
            size_t thrs, 
            Before_Exec_Task_Func func = Before_Exec_Task_Func{}
        );

        thread_pool(
            size_t thrs, 
            ::std::shared_ptr<TaskM> tasks_mgr, 
            Before_Exec_Task_Func func = Before_Exec_Task_Func{}
        );

        thread_pool(
            size_t thrs, 
            Sleeper slpr, 
            Before_Exec_Task_Func func = Before_Exec_Task_Func{}
        );

        thread_pool(
            size_t thrs, 
            ::std::shared_ptr<TaskM> tasks_mgr, 
            Sleeper slpr,
            Before_Exec_Task_Func func = Before_Exec_Task_Func{}
        );

        ~thread_pool();

    private:
        void start();
        void task_loop();

    public:
        void stop();
        void join();
        void blocked_stop();
        void schedule(resumable r) 
        { 
            task_manager().enqueue_task(::std::move(r)); 
            awake(); 
        }

        TaskM& task_manager() noexcept { return *tp_tasks; }
        void awake() noexcept { tp_sleeper.awake(); }
        auto& sleeper() noexcept { return tp_sleeper; }

    protected:
        bool done(::std::stop_token& stp) const noexcept 
        { 
            return stp.stop_requested() 
                && tp_tasks->empty() 
                && tp_working_thr_cnt.load() == 0 
                && tp_sleeper.done();
        }

        void execute_a_task(moodycamel::ConsumerToken& tk)
        {
            resumable cur{ tp_tasks->get_task(tk) };
            if (cur.empty()) tp_sleeper();
            else 
            {
                ++tp_working_thr_cnt;
                try { 
                    cur(); 
                } catch (...) { /* DO NOTHING*/ }
                --tp_working_thr_cnt;
            }
        }

    private:
        ::std::atomic_size_t tp_working_thr_cnt;
        sleeper_type tp_sleeper;
        ::std::stop_source tp_stop_src;
        const size_t tp_thrs_cnt{};
        ::std::shared_ptr<task_manager_type> 
            tp_tasks{ task_manager_type::get_from_global() };
        ::std::vector<::std::jthread> tp_threads;       
        Before_Exec_Task_Func tp_before_exec_task;
    };

    template<typename TP>
    concept is_thread_pool = util::is_specialization_of<TP, thread_pool>;

    template<task_manager TaskM, sleeper Sleeper, ::std::invocable Before_Exec_Task_Func>
    thread_pool<TaskM, Sleeper, Before_Exec_Task_Func>::
    thread_pool(size_t thrs, Before_Exec_Task_Func func)
        : tp_thrs_cnt{ thrs }, 
          tp_tasks{ TaskM::get_from_global() }, 
          tp_before_exec_task{ ::std::move(func) }
    {
        start();
    }

    template<task_manager TaskM, sleeper Sleeper, ::std::invocable Before_Exec_Task_Func>
    thread_pool<TaskM, Sleeper, Before_Exec_Task_Func>::
    thread_pool(size_t thrs, ::std::shared_ptr<TaskM> tasks, Before_Exec_Task_Func func)
        : tp_thrs_cnt{ thrs },
          tp_threads{ thrs }, 
          tp_before_exec_task{ ::std::move(func) }
    {
        start();
    }

    template<task_manager TaskM, sleeper Sleeper, ::std::invocable Before_Exec_Task_Func>
    thread_pool<TaskM, Sleeper, Before_Exec_Task_Func>::
    thread_pool(size_t thrs, Sleeper slpr, Before_Exec_Task_Func func)
        : tp_thrs_cnt{ thrs }, 
          tp_sleeper{ ::std::move(slpr) }, 
          tp_before_exec_task{ ::std::move(func) }
    {
        start();
    }

    template<task_manager TaskM, sleeper Sleeper, ::std::invocable Before_Exec_Task_Func>
    thread_pool<TaskM, Sleeper, Before_Exec_Task_Func>::
    thread_pool(size_t thrs, 
                ::std::shared_ptr<TaskM> tasks, 
                Sleeper slpr, 
                Before_Exec_Task_Func func)
        : tp_thrs_cnt{ thrs }, 
          tp_tasks{ tasks }, 
          tp_sleeper{ ::std::move(slpr) }, 
          tp_before_exec_task{ ::std::move(func) }
    {
        start();
    }

    template<task_manager TaskM, sleeper Sleeper, ::std::invocable Before_Exec_Task_Func>
    thread_pool<TaskM, Sleeper, Before_Exec_Task_Func>::~thread_pool()
    {
        blocked_stop();
    }

    template<task_manager TaskM, sleeper Sleeper, ::std::invocable Before_Exec_Task_Func>
    void thread_pool<TaskM, Sleeper, Before_Exec_Task_Func>::task_loop()
    {
        ::std::stop_token stp = tp_stop_src.get_token();
        auto consumer_tk = tp_tasks->get_moodycamel_consumer_tk();
        while (!done(stp))
        {
            tp_before_exec_task();
            execute_a_task(consumer_tk);
        }
    }

    template<task_manager TaskM, sleeper Sleeper, ::std::invocable Before_Exec_Task_Func>
    void thread_pool<TaskM, Sleeper, Before_Exec_Task_Func>::start()
    {
        for (size_t i = 0; i < tp_thrs_cnt; ++i)
        {
            tp_threads.emplace_back([this]{ task_loop(); });
        }
    }

    template<task_manager TaskM, sleeper Sleeper, ::std::invocable Before_Exec_Task_Func>
    void thread_pool<TaskM, Sleeper, Before_Exec_Task_Func>::stop()
    {
        tp_stop_src.request_stop();
    }

    template<task_manager TaskM, sleeper Sleeper, ::std::invocable Before_Exec_Task_Func>
    void thread_pool<TaskM, Sleeper, Before_Exec_Task_Func>::join()
    {
        auto thrs = ::std::move(tp_threads);
        for (auto& jt : thrs)
            jt.join();
    }

    template<task_manager TaskM, sleeper Sleeper, ::std::invocable Before_Exec_Task_Func>
    void thread_pool<TaskM, Sleeper, Before_Exec_Task_Func>::blocked_stop()
    {
        stop();
        join();
    }

    template<is_thread_pool Schr, typename Ivkb>
    void schedule(Schr& schr, Ivkb&& i) 
    {
        schr.schedule(resumable{ ::std::forward<Ivkb>(i) });
    }
}

#endif
